# 8个线性天线3层2次多项式拟合：12天线、sslDb=-30dB、扫描角度30度
# v5: 验证是否可以收敛于任意曲线
import argparse
import sys
from typing import Dict, Tuple
import numpy as np
from scipy.signal import windows as SSW # import chebwin
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from torch.optim.lr_scheduler import ReduceLROnPlateau
from apps.nlbf.nlbf_config import NlbfConfig as NG
from apps.nlbf.nlbf_model_v5 import NlbfModel


class NlbfApp(object):
    def __init__(self):
        self.name = 'apps.nlbf.nlbf_app.NlbfApp'

    @staticmethod
    def startup(params:Dict = {}) -> None:
        print(f'非线性波束成形应用 v5 001')
        # 设置中文字体
        plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
        plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题        
        if params['run_mode'] == 1: # 生成30度扫描角数据集
            NG.N = 8
            NG.rank = 3
            NG.layers = 3
            NlbfApp.generate_ds(
                X_fn = './work/nlbf/v5/a030/X_8.pt',
                y_fn = './work/nlbf/v5/a030/y_16_30dB_30.pt',
                N_base = 8,
                N_aim = 16,
                ssl_db = 30,
                theta = 30.0
            )
        elif params['run_mode'] == 2: # 模型训练
            NG.rank = 3 # 设置为2次多项式
            NG.layers = 3 # 3层网络
            NlbfApp.train_main(
                pt_fn='./work/nlbf/v5/a030/nlbf.pt', 
                X_fn='./work/nlbf/v5/a030/X_8.pt', 
                y_fn='./work/nlbf/v5/a030/y_16_30dB_30.pt'
            )
        elif params['run_mode'] == 3: # 模型训练效果可视化
            NG.rank = 3
            NG.layers = 3
            NlbfApp.evaluate(
                pt_fn='./work/nlbf/v5/a030/nlbf.pt', 
                X_fn='./work/nlbf/v5/a030/X_8.pt',  
                y_fn='./work/nlbf/v5/a030/y_16_30dB_30.pt'
            )
        elif params['run_mode'] == 1000: # 实例演示
            NG.rank = 3
            NG.layers = 3
            params['pt_fn'] = './work/nlbf/v5/a030/nlbf.pt'
            NlbfApp.mvdr_main(params=params)
        elif params['run_mode'] == 100: # 绘制理想图形
            NlbfApp.get_y_gt()
        elif params['run_mode'] == 999: # 临时试验程序
            NlbfApp.t001()
        else:
            print(f'未知运行模式')

    @staticmethod
    def t001():
        L = [8, 16, 32, 16, 8, 1]
        neurons_num = 0
        for lidx in range(len(L)-1):
            neurons_num += L[lidx]*L[lidx+1] + L[lidx+1]
        print(f'### neurons: {neurons_num};')

    @staticmethod
    def get_y_gt():
        '''
        生成带有干扰信号的目标波束成形结果
        X(1000,8) + scan_v(8,1) => MLP权重
        A(1800,8) => MLP => (1800,) 符合增强信号压制干扰的效果 batch_size=1800
        '''
        NG.theta = 30 # 扫描角度
        ssl_dB = -30.0 # 旁瓣水平
        NG.theta_rad = NG.theta / 180.0 * np.pi
        B = np.ones_like(NG.thetas_rad[0]) * ssl_dB
        print(f'### B: {B.shape};')
        # 处理信号
        sig_idx1, sig_idx2, sig_idx3 = NlbfApp.get_neighborhood(center=NG.theta_rad, radius=0.02*np.pi)
        NlbfApp.gen_plus(arrs=B, idx1=sig_idx1, idx2=sig_idx2, idx3=sig_idx3, v1=ssl_dB, v2=0.0, v3=ssl_dB)
        # 干扰处理
        intfrs_1_1, intfrs_1_2, intfrs_1_3 = NlbfApp.get_neighborhood(center=np.pi/3.0, radius=0.02*np.pi)
        intfrs_ssl_dB = -60.0
        NlbfApp.gen_plus(arrs=B, idx1=intfrs_1_1, idx2=intfrs_1_2, idx3=intfrs_1_3, v1=ssl_dB, v2=intfrs_ssl_dB, v3=ssl_dB)
        # plt.figure(figsize=(10, 6))
        # plt.plot(NG.thetas[0], B, label='MVDR Beam Pattern')
        # plt.show()
        B = 10.0**(B/20.0)
        return B

    @staticmethod
    def gen_plus(arrs:np.ndarray, 
        idx1:int, idx2:int, idx3:int, 
        v1:torch.float64, v2:torch.float64, v3:torch.float64
    ) -> None:
        slope1 = (v2-v1) / (idx2 - idx1)
        for idx in range(idx1, idx2+1, 1):
            # (x-ssl_dB) / (idx - sig_idx1) = slope1
            arrs[idx] = slope1 * (idx - idx1) + v1
        slope2 = (v3 - v2) / (idx3 - idx2)
        for idx in range(idx2, idx3, 1):
            # slope2 = (x - 0) / (idx - idx_2)
            arrs[idx] = slope2*(idx-idx2) + v2

    @staticmethod
    def get_neighborhood(center:torch.float64, radius:torch.float64) -> Tuple[int, int, int]:
        start_idx = 0
        end_idx = 0
        min_d1 = 1.0E8
        mid_idx = -1
        max_d2 = -1.0
        max_d2_idx = -1
        for idx, ri in enumerate(NG.thetas_rad[0]):
            d1 = center - ri
            d2 = ri - center
            if d1 > 0 and d1 < radius and start_idx == 0:
                start_idx = idx
                min_d1 = d1
            if d1 > 0 and d1 < radius and d1 < min_d1:
                min_d1 = d1
                mid_idx = idx
            if d2 > 0 and d2 < radius and d2 > max_d2:
                max_d2 = d2
                end_idx = idx
        return start_idx, mid_idx, end_idx
    
    @staticmethod
    def train_main(pt_fn:str = './work/nlbf/v5/a000/nlbf.pt', X_fn:str = './work/nlbf/v5/a000/X_nlbf.pt', y_fn:str = './work/nlbf/v5/a000/y_nlbf.pt') -> None:
        '''
        不能收敛的训练过程
        '''
        print(f'{NG.rank}阶多项式回归 v0.0.1')
        NG.theta = 30
        NG.theta_rad = NG.theta / 180 * np.pi
        N_aim = 8


        NG.N = 8
        NG.rank = 3
        NG.layers = 3
        NG.theta = 30
        N_aim = 8
        NG.theta_rad = NG.theta / 180.0 * np.pi
        theta0 = NG.theta # 30    # 期望信号方向（度）
        theta1 = 60    # 干扰信号方向（度）
        signal_power = 1
        interference_power = 0.5
        noise_power = 0.1
        theta0_rad = NlbfApp.steering_vector(theta0, NG.N, d=NG.d[0], wavelength=NG.lambda_) # (8,)
        theta1_rad = NlbfApp.steering_vector(theta1, NG.N, d=NG.d[0], wavelength=NG.lambda_) # (8,)
        # 生成信号
        s = np.sqrt(signal_power) * np.exp(1j*2*np.pi*(NG.f*NG.times + 0.5*NG.slope*NG.times**2))
        i = np.sqrt(interference_power) * (np.random.randn(NG.K) + 1j * np.random.randn(NG.K)) # (1000,)
        n = np.sqrt(noise_power) * (np.random.randn(NG.N, NG.K) + 1j * np.random.randn(NG.N, NG.K)) # (8, 1000)
        # 接收信号矩阵
        X_np = np.outer(theta0_rad, s) + np.outer(theta1_rad, i) + n # (8, 1000)
        X = torch.from_numpy(X_np.T)
        X = X.to(torch.complex128).to(NG.device)





        # 载入波束成形相关内容和训练目标
        A = torch.load(X_fn, weights_only=True) # 导向向量 X: ((8,1800), complex128)
        A = A.T
        A = A.to(NG.device) # X: ((1800,8), complex128)
        y = torch.load(y_fn, weights_only=True).unsqueeze(1) # y: ((1800,1), float64)
        y = y.to(NG.device)
        y0_np = NlbfApp.get_y_gt()
        y0 = torch.from_numpy(y0_np)
        y0 = y0.unsqueeze(1)
        y0 = y0.to(torch.float64).to(NG.device)
        print(f'### y: {y[0][0]}; {y[1][0]}; {y[2][0]}; {y[3][0]}; {y[4][0]};')
        print(f'### y0: {y0[0][0]}; {y0[1][0]}; {y0[2][0]}; {y0[3][0]}; {y0[4][0]};')
        theta = NG.theta
        scan_raw = np.exp(1j * 2 * np.pi / NG.lambda_ * np.arange(N_aim) * NG.d[0] * np.sin(NG.theta_rad))
        scan_w = torch.from_numpy(scan_raw)
        scan_w = scan_w.to(torch.complex128)
        scan_w = scan_w.to(NG.device)

        # 创建模型实例
        print(f'### X: {A.shape}; y: {y.shape};')
        model = NlbfModel(avec_dim=1, rank=NG.rank, layers=NG.layers)
        model = model.to(NG.device)
        # 定义损失函数
        criterion = nn.MSELoss()
        # criterion = nn.L1Loss()
        # 定义优化器
        optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)
        # 预热阶段
        print(f'##### 开始预热训练')
        warmup_opt = torch.optim.AdamW(model.parameters(), lr=1e-5, betas=(0.9, 0.99), weight_decay=1e-2)
        warmup_scheduler = ReduceLROnPlateau(warmup_opt, 'min', threshold=1e-4, eps=1e-8)
        warmup_epochs = 5
        for epoch in range(warmup_epochs):
            print(f'    warmup {epoch}......')
            NlbfApp.train_epoch(epoch, model, criterion, warmup_opt, warmup_scheduler, X, scan_w, A, y0)
        # 正式训练过程
        print(f'##### 开始正式训练')
        optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, betas=(0.9, 0.99), weight_decay=1e-2)
        scheduler = ReduceLROnPlateau(optimizer, 'min', threshold=1e-4, eps=1e-8)
        epochs = 2000000
        best_loss = sys.float_info.max
        improve_threshold = 1e-7
        cumulative_steps = 0
        max_unimproves = 500
        for epoch in range(epochs):
            loss = NlbfApp.train_epoch(epoch, model, criterion, optimizer, scheduler, X, scan_w, A, y0)
            print(f'epoch_{epoch}: loss={loss.item()};')
            cumulative_steps += 1
            if best_loss > loss:
                if best_loss - loss > improve_threshold:
                    print(f'Save best weights')
                    model.save(pt_fn=pt_fn)
                    cumulative_steps = 0
                    best_loss = loss
            if cumulative_steps > max_unimproves:
                print(f'Earlly Stopping!!!!!!')
                break

    @staticmethod
    def train_epoch(epoch, model, criterion, optimizer, scheduler, X, scan_w, A, y_gt):
        batch_size = 10
        max_loss = torch.tensor([0.0]).to(X.device)
        total_loss = torch.tensor([0.0]).to(X.device)
        # 训练过程
        for bidx in range(batch_size):
            loss = None
            y_hat = None
            y_hat = model(X, scan_w, A)
            loss = criterion(y_hat, y_gt)
            if loss > max_loss:
                max_loss = loss
            total_loss += loss
            # 清空梯度
            optimizer.zero_grad()
            # 反向传播
            loss.backward()
            # 更新参数
            optimizer.step()
        if scheduler is not None:
            scheduler.step(total_loss)
        return max_loss
            

    @staticmethod
    def evaluate(pt_fn:str = './work/nlbf/v5/nlbf.pt', X_fn:str = './work/nlbf/v5/X_nlbf.pt', y_fn:str = './work/nlbf/v5/y_nlbf.pt') -> None:
        print(f'非线性波束形成简单评估({NG.layers}层{NG.rank-1}次多项式回归) v0.0.1 001')


        NG.N = 8
        NG.rank = 3
        NG.layers = 3
        NG.theta = 30
        N_aim = 8
        NG.theta_rad = NG.theta / 180.0 * np.pi
        theta0 = NG.theta # 30    # 期望信号方向（度）
        theta1 = 60    # 干扰信号方向（度）
        signal_power = 1
        interference_power = 0.5
        noise_power = 0.1
        theta0_rad = NlbfApp.steering_vector(theta0, NG.N, d=NG.d[0], wavelength=NG.lambda_) # (8,)
        theta1_rad = NlbfApp.steering_vector(theta1, NG.N, d=NG.d[0], wavelength=NG.lambda_) # (8,)
        # 生成信号
        s = np.sqrt(signal_power) * np.exp(1j*2*np.pi*(NG.f*NG.times + 0.5*NG.slope*NG.times**2))
        i = np.sqrt(interference_power) * (np.random.randn(NG.K) + 1j * np.random.randn(NG.K)) # (1000,)
        n = np.sqrt(noise_power) * (np.random.randn(NG.N, NG.K) + 1j * np.random.randn(NG.N, NG.K)) # (8, 1000)
        # 接收信号矩阵
        X_np = np.outer(theta0_rad, s) + np.outer(theta1_rad, i) + n # (8, 1000)
        X = torch.from_numpy(X_np.T)
        X = X.to(torch.complex128).to(NG.device)




        NG.theta = 30
        NG.theta_rad = NG.theta / 180 * np.pi
        NG.N = 8
        NG.rank = 3
        NG.layers = 3
        n_arr = np.arange(NG.N).reshape((-1, 1))
        N_aim = 8
        # scan_raw = np.exp(1j * 2 * np.pi / NG.lambda_ * np.arange(N_aim) * NG.d[0] * np.sin(NG.theta_rad))
        scan_raw = np.exp(1j * 2 * np.pi * np.sin(NG.theta_rad) * n_arr * NG.d[0] / NG.lambda_)
        scan_w = torch.from_numpy(scan_raw)
        scan_w = scan_w.to(torch.complex128)
        scan_w = scan_w.to(NG.device)
        print(f'scan_w:\n{scan_w};\n')
        # Tschebyscheff波束形成
        theta = NG.theta # 扫描角度
        theta_rad = theta / 180 * np.pi
        a = np.exp(1j * 2 * np.pi * np.sin(NG.thetas_rad) * n_arr * NG.d[0] / NG.lambda_)
        w = np.exp(1j * 2 * np.pi * np.sin(theta_rad) * n_arr * NG.d[0] / NG.lambda_)
        B = np.dot(a.transpose(), np.conj(w))
        B = np.abs(B) / np.max(np.abs(B))
        # 载入波束成形相关内容和训练目标
        A = torch.load(X_fn, weights_only=True) # 导向向量
        A = A.T
        A = A.to(NG.device)
        y = torch.load(y_fn, weights_only=True).unsqueeze(1)
        y = y.to(NG.device)

        y0_np = NlbfApp.get_y_gt()
        y0 = torch.from_numpy(y0_np)
        y0 = y0.unsqueeze(1)
        y0 = y0.to(torch.float64).to(NG.device)



        print(f'### A: {A.shape}, {A.dtype}; ?????????????')
        print(f'### y: {y.shape}, {y.dtype}; {y[0]};')
        # 创建模型实例
        model = NlbfModel(avec_dim=1, rank=NG.rank, layers=NG.layers)
        model.load(pt_fn = pt_fn) # 载入权重文件
        model = model.to(NG.device)
        y_hat = model(X, scan_w.squeeze(1), A)
        torch.save(X, './work/nlbf/v5/X.pt')
        torch.save(scan_w, './work/nlbf/v5/sw.pt')
        # 进行可视化
        # 训练后的结果评估
        a_np = X.cpu().numpy()
        # 天线图响应
        y_np = y.detach().cpu().numpy()
        y_hat_np = y_hat.detach().cpu().numpy()
        print(f'{y_hat_np}')
        # y_hat_np = np.abs(y_hat_np) / np.max(np.abs(y_hat_np))
        # 绘制波束形成图
        fig, ax = plt.subplots(1, 1, sharex=True, sharey=True, figsize=(7, 2), dpi=300)
        plt.subplots_adjust(wspace=0.1, hspace=0.05)
        ax.plot(NG.thetas[0], 20 * np.log10(y0_np), '--', label='16天线')
        ax.plot(NG.thetas[0], 20 * np.log10(y_hat_np), '', label=f'{NG.layers}层{NG.rank-1}次多项式')
        ax.plot(NG.thetas[0], 20 * np.log10(B), '--', label='8天线')
        ax.legend()
        ax.set_title(f'波束成形结果({NG.layers}层{NG.rank-1}次)', fontsize=12, color='blue', loc='center')
        ax.set_ylim(-60, 1)
        # plt.show()
        plt.savefig('./work/nlbf/v5/a001.png')

    @staticmethod
    def generate_ds(X_fn:str = './work/nlbf/v5/a000/X_nlbf.pt', y_fn:str = './work/nlbf/v5/a000/y_nlbf.pt', N_base:int = 8, N_aim:int = 32, ssl_db:float = -30, theta:float = 0.0) -> None:
        '''
        N_base: 基础天线数量
        N_target: 优化目标的天线数量
        ssl_db: 旁瓣水平
        '''
        theta_rad = theta / 180.0 * np.pi
        af = SSW.chebwin(N_aim, at=ssl_db, sym=True) # af: (16,) float64
        phases = np.exp(1j * 2 * np.pi / NG.lambda_ * np.arange(N_aim) * NG.d[0] * np.sin(theta_rad)) # phases: (16,) complex128
        w_vec_aim = af * phases # 权重向量 (16,) complex128
        # 导向向量 a_vec_aim: ((16, 1800), complex128)
        a_vec_aim = np.array([np.exp(1j * 2 * np.pi / NG.lambda_ * np.arange(N_aim) * NG.d[0] * np.sin((ti))) for ti in NG.thetas_rad[0]]).T
        # 天线图响应 B: ((1800,), float64)
        B = np.dot(a_vec_aim.transpose(), np.conj(w_vec_aim))
        B = np.abs(B) / np.max(np.abs(B))
        # a_vec_base: ((8,1800), complex128)
        a_vec_base = np.array([np.exp(1j * 2 * np.pi / NG.lambda_ * np.arange(N_base) * NG.d[0] * np.sin((ti))) for ti in NG.thetas_rad[0]]).T
        X = torch.from_numpy(a_vec_base)
        torch.save(X, X_fn)
        y = torch.from_numpy(B) # ((1800,), float64)
        torch.save(y, y_fn)

    

    @staticmethod
    def mvdr_main(params:Dict = {}) -> None:
        '''
        最小方差无失真响应（MVDR）
        '''
        NG.N = 8
        NG.rank = 3
        NG.layers = 3
        NG.theta = 30
        N_aim = 8
        NG.theta_rad = NG.theta / 180.0 * np.pi
        pt_fn = params['pt_fn']
        theta0 = NG.theta # 30    # 期望信号方向（度）
        theta1 = 60    # 干扰信号方向（度）
        signal_power = 1
        interference_power = 0.5
        noise_power = 0.1
        theta0_rad = NlbfApp.steering_vector(theta0, NG.N, d=NG.d[0], wavelength=NG.lambda_) # (8,)
        theta1_rad = NlbfApp.steering_vector(theta1, NG.N, d=NG.d[0], wavelength=NG.lambda_) # (8,)
        # 生成信号
        s = np.sqrt(signal_power) * np.exp(1j*2*np.pi*(NG.f*NG.times + 0.5*NG.slope*NG.times**2))
        i = np.sqrt(interference_power) * (np.random.randn(NG.K) + 1j * np.random.randn(NG.K)) # (1000,)
        n = np.sqrt(noise_power) * (np.random.randn(NG.N, NG.K) + 1j * np.random.randn(NG.N, NG.K)) # (8, 1000)
        # 接收信号矩阵
        X = np.outer(theta0_rad, s) + np.outer(theta1_rad, i) + n # (8, 1000)
        # 计算MVDR权重
        w = NlbfApp.mvdr_beamforming(X, theta0_rad) # (8,)
        # 绘制波束形成图
        # 波束形成
        ##########################
        avb_np = np.array([np.exp(1j * 2 * np.pi / NG.lambda_ * np.arange(NG.N) * NG.d[0] * np.sin((ti))) for ti in NG.thetas_rad[0]]).T
        avb_tr = torch.from_numpy(avb_np)
        avb_tr = avb_tr.to(torch.complex128).to(NG.device)
        # # 线性优化
        B = np.dot(avb_np.T, np.conj(w)) # avb_np: (1800, 8); w: (8,) => B: (1800,)
        B = np.abs(B) / np.max(np.abs(B))
        ################################################################################################################
        # ????? 使用训练好的模型，查看效果
        NG.theta = 30
        NG.theta_rad = NG.theta / 180.0 * np.pi
        theta = NG.theta
        n_arr = np.arange(NG.N).reshape((-1, 1))
        scan_raw = np.exp(1j * 2 * np.pi * np.sin(NG.theta_rad) * n_arr * NG.d[0] / NG.lambda_)
        scan_w = torch.from_numpy(scan_raw)
        scan_w = scan_w.to(torch.complex128).to(NG.device)
        # 创建模型实例
        model = NlbfModel(avec_dim=1, rank=NG.rank, layers=NG.layers)
        model.load(pt_fn = pt_fn) # 载入权重文件
        model = model.to(NG.device)
        y_hat = model(avb_tr.T, scan_w.squeeze(1))
        y_hat_np = y_hat.detach().cpu().numpy()
        ################################################################################################################
        # 非线性优化
        fig, ax = plt.subplots(1, 2, sharex=True, sharey=True, figsize=(7, 2), dpi=300)
        plt.subplots_adjust(wspace=0.1, hspace=0.05)
        ax[0].plot(NG.thetas[0], 20 * np.log10(B), '', color='r')
        ax[1].plot(NG.thetas[0], 20 * np.log10(y_hat_np), '', color='g')
        plt.show()
        # plt.savefig('./work/nlbf/v5/a002.png')
        # 输出信号
        y = np.dot(w.conj().T, X) # (1000,)
        # 绘制信号功率谱
        plt.figure(figsize=(10, 6))
        # angles = np.linspace(0, 75, 181)
        responses = []
        for theta in NG.thetas[0]:
            a_theta = NlbfApp.steering_vector(theta, NG.N, d=NG.d[0], wavelength=NG.lambda_)
            response = np.abs(np.dot(w.conj().T, a_theta))**2
            responses.append(response)
        z = model(avb_tr.T, scan_w.squeeze(1), mode=2)
        z_np = z.detach().cpu().numpy().squeeze(0)
        plt.plot(NG.thetas[0], 10 * np.log10(responses / np.max(responses)), label='MVDR Beam Pattern')
        plt.plot(NG.thetas[0], 10 * np.log10(z_np), label='NLBF Beam Pattern')
        plt.title('MVDR波束图 v5')
        plt.xlabel('方向角度（度）')
        plt.ylabel('归一化增益（dB）')
        plt.grid(True)
        plt.legend()
        plt.show()
        # plt.savefig('./work/nlbf/v5/a003.png')
        # 输出权重向量
        print("MVDR权重向量:")
        print(w)


    # 生成阵列响应向量
    def steering_vector(theta, N, d=0.5, wavelength=1):
        angles = np.deg2rad(theta)
        return np.exp(1j * 2 * np.pi * d * np.arange(N) * np.sin(angles) / wavelength)
        
    def mvdr_beamforming(X, a_theta0):
        """
        MVDR波束成形
        X: 接收信号矩阵 (N x K)
        a_theta0: 期望方向的阵列响应向量 (N x 1)
        """
        R = np.dot(X, X.conj().T) / X.shape[1]  # 计算协方差矩阵
        R_inv = np.linalg.inv(R)
        numerator = np.dot(R_inv, a_theta0)
        denominator = np.dot(a_theta0.conj().T, numerator)
        w_mvdr = numerator / denominator
        return w_mvdr










def main(params:Dict = {}) -> None:
    NlbfApp.startup(params=params)

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--run_mode', action='store',
        type=int, default=1, dest='run_mode',
        help='run mode'
    )
    return parser.parse_args()

if '__main__' == __name__:
    args = parse_args()
    params = vars(args)
    main(params=params)